package com.jh.cloud.servlet.gateway.servlet;

import com.jh.cloud.servlet.gateway.loadbalance.ZookeeperLoadBalancer;
import com.netflix.loadbalancer.Server;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.RequestEntity;
import org.springframework.http.ResponseEntity;
import org.springframework.util.MultiValueMap;
import org.springframework.util.StreamUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.client.RestTemplate;

import javax.servlet.ServletOutputStream;
import javax.servlet.annotation.WebServlet;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.InputStream;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.Collections;
import java.util.List;
import java.util.Optional;

@WebServlet(name = "loadbalance", urlPatterns = {"/loadbalance/*"})
@SuppressWarnings("all")
public class LoadBalanceGatewayServlet extends HttpServlet {

    @Autowired
    private ZookeeperLoadBalancer zookeeperLoadBalancer;

    /**
     * 选择一个实例
     *
     * @param serviceName 服务名称
     * @return 服务
     */
    public Server chooseServer(String serviceName) {
        return zookeeperLoadBalancer.chooseServer(serviceName);
    }

    /**
     * 构建目标URL
     *
     * @param serviceInstance 服务实例
     * @param serviceUri      服务uri
     * @param req             请求
     * @return 目标URL
     */
    private String buildTargetURL(Server serviceInstance, String serviceUri, HttpServletRequest req) {
        StringBuilder urlBuilder = new StringBuilder();
        urlBuilder.append(serviceInstance.getScheme())
                .append(serviceInstance.getHost())
                .append(":")
                .append(serviceInstance.getPort())
                .append(serviceUri);
        String queryString = req.getQueryString();
        if (StringUtils.hasLength(queryString)) {
            urlBuilder.append("?").append(queryString);
        }

        return urlBuilder.toString();
    }

    /**
     * 获取 request body
     *
     * @param request 请求
     * @return request body
     * @throws IOException IO异常
     */
    private byte[] createRequestBody(HttpServletRequest request) throws IOException {
        InputStream inputStream = request.getInputStream();
        return StreamUtils.copyToByteArray(inputStream);
    }

    /**
     * 获取 headers
     *
     * @param request 请求
     * @return header map
     */
    private MultiValueMap<String, String> createRequestHeaders(HttpServletRequest request) {
        HttpHeaders headers = new HttpHeaders();
        List<String> headerNames = Collections.list(request.getHeaderNames());
        headerNames.forEach(headerName -> {
            List<String> headerValues = Collections.list(request.getHeaders(headerName));
            headerValues.forEach(headerValue -> headers.add(headerName, headerValue));
        });
        return headers;
    }

    /**
     * 创建 request entity
     *
     * @param request 请求
     * @param uri     uri
     * @return request entity
     * @throws IOException        IO异常
     * @throws URISyntaxException URI转换异常
     */
    private RequestEntity<byte[]> createRequestEntity(HttpServletRequest request, String uri) throws IOException, URISyntaxException {
        // method string
        String method = request.getMethod();
        // wrap method to http-method
        HttpMethod httpMethod = HttpMethod.resolve(method);
        // get request body
        byte[] body = createRequestBody(request);
        // get headers
        MultiValueMap<String, String> headers = createRequestHeaders(request);
        // request entity
        return new RequestEntity<>(body, headers, httpMethod, new URI(uri));
    }

    /**
     * 将 RestTemplate 的响应的Header信息输出到 HttpResponse
     *
     * @param responseEntity 响应体
     * @param resp           HttpServletResponse
     */
    private void writeHeaders(ResponseEntity<byte[]> responseEntity, HttpServletResponse resp) {
        // get http headers
        HttpHeaders httpHeaders = responseEntity.getHeaders();
        // write for response
        httpHeaders.forEach((headerName, headerValues) -> headerValues.forEach(headerValue -> resp.addHeader(headerName, headerValue)));
    }

    /**
     * 将 RestTemplate 的响应body输出到 HttpServletResponse
     *
     * @param responseEntity 响应体
     * @param resp           HttpServeletResponse
     * @throws IOException IO异常
     */
    private void writeBody(ResponseEntity<byte[]> responseEntity, HttpServletResponse resp) throws IOException {
        if (responseEntity.hasBody()) {
            byte[] body = responseEntity.getBody();
            // 获取输出流
            ServletOutputStream output = resp.getOutputStream();
            // 将 RestTemplate 响应内容写到 Http 输出流
            output.write(body);
            output.flush();
        }
    }

    @Override
    protected void service(HttpServletRequest req, HttpServletResponse resp) throws IOException {

        // {service-name}/{service-uri}
        String pathInfo = req.getPathInfo().substring(1);

        String[] parts = Optional.ofNullable(StringUtils.split(pathInfo, "/")).orElseThrow(() -> new RuntimeException("parts is null"));

        // service-name
        String serviceName = parts[0];
        // service-uri
        String serviceUri = "/" + parts[1];
        // service-instance
        Server serviceInstance = Optional.ofNullable(chooseServer(serviceName)).orElseThrow(() -> new RuntimeException("service instance is null"));
        // 构建URL：scheme://ip:port/service-uri
        String targetUrl = buildTargetURL(serviceInstance, serviceUri, req);

        // 转发
        RestTemplate restTemplate = new RestTemplate();

        // 构造 request entity
        RequestEntity<byte[]> requestEntity = null;
        try {
            requestEntity = createRequestEntity(req, targetUrl);
            ResponseEntity<byte[]> responseEntity = restTemplate.exchange(requestEntity, byte[].class);
            writeHeaders(responseEntity, resp);
            writeBody(responseEntity, resp);
        } catch (URISyntaxException e) {
            e.printStackTrace();
        }
    }

}
